"""
Demo 01: 原始数值 + 基础预处理（窗口化）
使用AAPL股票收盘价，进行对数收益计算、Z-score标准化和滑动窗口处理
"""

import yfinance as yf
import pandas as pd
import numpy as np


class RawVectorizer:
    """原始数值窗口化向量处理类"""
    
    def __init__(self, window_size=60, step_size=5):
        """
        初始化参数
        
        Args:
            window_size: 窗口大小
            step_size: 步长
        """
        self.window_size = window_size
        self.step_size = step_size
        self.mean = None
        self.std = None
        
    def download_data(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        下载股票数据
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            pandas.Series: 收盘价序列
        """
        df = yf.download(ticker, start=start, end=end, progress=False)
        return df["Close"]
    
    def compute_returns(self, price_series):
        """
        计算对数收益率
        
        Args:
            price_series: 价格序列
            
        Returns:
            pandas.Series: 对数收益率序列
        """
        returns = np.log(price_series).diff().dropna()
        return returns
    
    def z_score_normalize(self, returns):
        """
        Z-score标准化
        
        Args:
            returns: 收益率序列
            
        Returns:
            pandas.Series: 标准化后的序列
        """
        self.mean = returns.mean()
        self.std = returns.std()
        z_returns = (returns - self.mean) / self.std
        return z_returns
    
    def create_windows(self, series):
        """
        创建滑动窗口
        
        Args:
            series: 时间序列
            
        Returns:
            numpy.ndarray: 窗口矩阵 (n_windows, window_size)
        """
        arr = series.values
        n_windows = (len(arr) - self.window_size) // self.step_size + 1
        windows = np.stack([
            arr[i * self.step_size : i * self.step_size + self.window_size] 
            for i in range(n_windows)
        ])
        return windows
    
    def fit_transform(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        完整的向量化流程
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            numpy.ndarray: 窗口化的向量矩阵
        """
        # 下载数据
        price = self.download_data(ticker, start, end)
        
        # 计算对数收益
        returns = self.compute_returns(price)
        
        # Z-score标准化
        z_returns = self.z_score_normalize(returns)
        
        # 创建窗口
        windows = self.create_windows(z_returns)
        
        return windows
    
    def transform(self, series):
        """
        对新数据进行转换（使用已有的均值和标准差）
        
        Args:
            series: 新的时间序列
            
        Returns:
            numpy.ndarray: 窗口化的向量矩阵
        """
        if self.mean is None or self.std is None:
            raise ValueError("Model not fitted yet. Call fit_transform first.")
        
        # 标准化
        z_series = (series - self.mean) / self.std
        
        # 创建窗口
        windows = self.create_windows(z_series)
        
        return windows


def main():
    """主函数，演示使用方法"""
    print("=" * 60)
    print("Demo 01: 原始数值 + 基础预处理（窗口化）")
    print("=" * 60)
    
    # 创建向量化器
    vectorizer = RawVectorizer(window_size=60, step_size=5)
    
    # 执行向量化
    print("正在下载AAPL股票数据...")
    X_raw = vectorizer.fit_transform(
        ticker="AAPL", 
        start="2020-01-01", 
        end="2024-12-31"
    )
    
    # 输出结果
    print(f"\n原始窗口 shape: {X_raw.shape}")
    print(f"窗口数量: {X_raw.shape[0]}")
    print(f"窗口长度: {X_raw.shape[1]}")
    
    # 显示第一个窗口的统计信息
    print(f"\n第一个窗口的统计信息:")
    print(f"  均值: {X_raw[0].mean():.4f}")
    print(f"  标准差: {X_raw[0].std():.4f}")
    print(f"  最小值: {X_raw[0].min():.4f}")
    print(f"  最大值: {X_raw[0].max():.4f}")
    
    # 返回结果供其他模块使用
    return X_raw


if __name__ == "__main__":
    X_raw = main()
